package com.zzl.study.sharding.algorithm.complex;

import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;

import java.math.BigInteger;
import java.util.*;

/**
 * @author: zhangzl
 * @date: 2022/7/17 12:00
 * @version: 1.0
 * @description: 书本表自定义类分片算法：complex策略，对书本表的adminId和region进行复合分片，不考虑范围分片（复合键查询的话就是全分片路由）
 * A区的放在0和1分片中，再根据adminId判断奇偶，偶数放在0分片，奇数放在1分片；
 * B区的放在2和3分片中，再根据adminId判断奇偶，偶数放在2分片，奇数放在3分片；
 * C区的放在4和5分片中，再根据adminId判断奇偶，偶数放在4分片，奇数放在5分片。
 */
public class BookComplexKeyShardingAlgorithm implements ComplexKeysShardingAlgorithm<Comparable<?>> {

    /**
     * 配置的props属性
     */
    private Properties props;

    /**
     * 存放分区对应的分片下标的Map集合
     */
    private Map<String, int[]> complexKeyMap = new HashMap<>();

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Comparable<?>> shardingValue) {
        List<String> result = new ArrayList<>();
        String logicTableName = shardingValue.getLogicTableName();
        Map<String, Collection<Comparable<?>>> columnNameAndShardingValuesMap = shardingValue.getColumnNameAndShardingValuesMap();
//        Map<String, Range<Comparable<?>>> columnNameAndRangeValuesMap = shardingValue.getColumnNameAndRangeValuesMap();
        Collection<Comparable<?>> regions = columnNameAndShardingValuesMap.get("region");
        Collection<Comparable<?>> adminIds = columnNameAndShardingValuesMap.get("admin_id");
        for (Comparable<?> region : regions) {
            String regionValue = (String) region;
            int[] suffixArray = this.complexKeyMap.get(regionValue);
            for (Comparable<?> adminId : adminIds) {
                Long adminIdValue = (Long) adminId;
                int index = new BigInteger(String.valueOf(adminIdValue)).mod(new BigInteger("2")).intValue();
                result.add(logicTableName + "_" + suffixArray[index]);
            }
        }
        return result;
    }

    @Override
    public Properties getProps() {
        return this.props;
    }

    @Override
    public void init(Properties props) {
        this.props = props;
        complexKeyMap.put("A区", new int[]{0, 1});
        complexKeyMap.put("B区", new int[]{2, 3});
        complexKeyMap.put("C区", new int[]{4, 5});
    }

}
